package com.newbee.utils;

import cn.hutool.core.util.IdUtil;
import com.newbee.model.TraceContext;
import org.apache.commons.lang3.StringUtils;


/**
 * 链路上下文工具类：管理traceId与节点链状态，处理线程安全
 */
public class TraceContextHolder {

    // 线程私有上下文，存储当前链路的traceId与节点信息
    private static final ThreadLocal<TraceContext> TRACE_CONTEXT = new ThreadLocal<>();

    /**
     * 获取或初始化上下文（自动生成/复用traceId）
     */
    public static TraceContext getOrInit(boolean init) {
        TraceContext context = TRACE_CONTEXT.get();
        if (context == null && init) {
            String traceId = StringUtils.defaultIfBlank(
                    HttpTraceHeaderUtil.getTraceIdFromHeader(),
                    IdUtil.fastSimpleUUID()
            );
            context = new TraceContext(traceId);
            TRACE_CONTEXT.set(context);
        }
        return context;
    }

    public static void initForAsync(String traceId, String asyncRootNodeId) {
        TraceContext context = new TraceContext(traceId, asyncRootNodeId);
        TRACE_CONTEXT.set(context);
    }

    public static ThreadLocal<TraceContext> getTraceContext() {
        return TRACE_CONTEXT;
    }

    /**
     * 清空上下文（必须在链路结束时调用，避免线程复用污染）
     */
    public static void clear() {
        TRACE_CONTEXT.remove();
    }
}